import paddle
"""
TensorFlow, Keras and TFLite versions of YOLOv5
Authored by https://github.com/zldrobit in PR https://github.com/ultralytics/yolov5/pull/1127

Usage:
    $ python models/tf.py --weights yolov5s.pt

Export:
    $ python path/to/export.py --weights yolov5s.pt --include saved_model pb tflite tfjs
"""
import argparse
import logging
import sys
from copy import deepcopy
from pathlib import Path
FILE = Path(__file__).resolve()
ROOT = FILE.parents[1]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))
import numpy as np
import tensorflow as tf
from tensorflow import keras
from models.common import Bottleneck, BottleneckCSP, Concat, Conv, C3, DWConv, Focus, SPP, SPPF, autopad
from models.experimental import CrossConv, MixConv2d, attempt_load
from models.yolo import Detect
from utils.general import make_divisible, print_args, set_logging
from utils.activations import SiLU
LOGGER = logging.getLogger(__name__)


class TFBN(keras.layers.Layer):

    def __init__(self, w=None):
        super(TFBN, self).__init__()
        self.bn = keras.layers.BatchNormalization(beta_initializer=keras.
            initializers.Constant(w.bias.numpy()), gamma_initializer=keras.
            initializers.Constant(w.weight.numpy()),
            moving_mean_initializer=keras.initializers.Constant(w.
            running_mean.numpy()), moving_variance_initializer=keras.
            initializers.Constant(w.running_var.numpy()), epsilon=w.eps)

    def call(self, inputs):
        return self.bn(inputs)


class TFPad(keras.layers.Layer):

    def __init__(self, pad):
        super(TFPad, self).__init__()
        self.pad = tf.constant([[0, 0], [pad, pad], [pad, pad], [0, 0]])

    def call(self, inputs):
        return tf.pad(inputs, self.pad, mode='constant', constant_values=0)


class TFConv(keras.layers.Layer):

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None):
        super(TFConv, self).__init__()
        assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument"
        assert isinstance(k, int
            ), 'Convolution with multiple kernels are not allowed.'
        conv = keras.layers.Conv2D(c2, k, s, 'SAME' if s == 1 else 'VALID',
            use_bias=False if hasattr(w, 'bn') else True,
            kernel_initializer=keras.initializers.Constant(w.conv.weight.
            transpose(perm=[2, 3, 1, 0]).numpy()), bias_initializer='zeros' if
            hasattr(w, 'bn') else keras.initializers.Constant(w.conv.bias.
            numpy()))
        self.conv = conv if s == 1 else keras.Sequential([TFPad(autopad(k,
            p)), conv])
        self.bn = TFBN(w.bn) if hasattr(w, 'bn') else tf.identity
        if isinstance(w.act, paddle.nn.LeakyReLU):
            self.act = (lambda x: keras.activations.relu(x, alpha=0.1)
                ) if act else tf.identity
        elif isinstance(w.act, paddle.nn.Hardswish):
            self.act = (lambda x: x * tf.nn.relu6(x + 3) * 0.166666667
                ) if act else tf.identity
        elif isinstance(w.act, (paddle.nn.Silu, SiLU)):
            self.act = (lambda x: keras.activations.swish(x)
                ) if act else tf.identity
        else:
            raise Exception(
                f'no matching TensorFlow activation found for {w.act}')

    def call(self, inputs):
        return self.act(self.bn(self.conv(inputs)))


class TFFocus(keras.layers.Layer):

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None):
        super(TFFocus, self).__init__()
        self.conv = TFConv(c1 * 4, c2, k, s, p, g, act, w.conv)

    def call(self, inputs):
        return self.conv(tf.concat([inputs[:, ::2, ::2, :], inputs[:, 1::2,
            ::2, :], inputs[:, ::2, 1::2, :], inputs[:, 1::2, 1::2, :]], 3))


class TFBottleneck(keras.layers.Layer):

    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5, w=None):
        super(TFBottleneck, self).__init__()
        c_ = int(c2 * e)
        self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
        self.cv2 = TFConv(c_, c2, 3, 1, g=g, w=w.cv2)
        self.add = shortcut and c1 == c2

    def call(self, inputs):
        return inputs + self.cv2(self.cv1(inputs)) if self.add else self.cv2(
            self.cv1(inputs))


class TFConv2d(keras.layers.Layer):

    def __init__(self, c1, c2, k, s=1, g=1, bias=True, w=None):
        super(TFConv2d, self).__init__()
        assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument"
        self.conv = keras.layers.Conv2D(c2, k, s, 'VALID', use_bias=bias,
            kernel_initializer=keras.initializers.Constant(w.weight.
            transpose(perm=[2, 3, 1, 0]).numpy()), bias_initializer=keras.
            initializers.Constant(w.bias.numpy()) if bias else None)

    def call(self, inputs):
        return self.conv(inputs)


class TFBottleneckCSP(keras.layers.Layer):

    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None):
        super(TFBottleneckCSP, self).__init__()
        c_ = int(c2 * e)
        self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
        self.cv2 = TFConv2d(c1, c_, 1, 1, bias=False, w=w.cv2)
        self.cv3 = TFConv2d(c_, c_, 1, 1, bias=False, w=w.cv3)
        self.cv4 = TFConv(2 * c_, c2, 1, 1, w=w.cv4)
        self.bn = TFBN(w.bn)
        self.act = lambda x: keras.activations.relu(x, alpha=0.1)
        self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0,
            w=w.m[j]) for j in range(n)])

    def call(self, inputs):
        y1 = self.cv3(self.m(self.cv1(inputs)))
        y2 = self.cv2(inputs)
        return self.cv4(self.act(self.bn(tf.concat((y1, y2), axis=3))))


class TFC3(keras.layers.Layer):

    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, w=None):
        super(TFC3, self).__init__()
        c_ = int(c2 * e)
        self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
        self.cv2 = TFConv(c1, c_, 1, 1, w=w.cv2)
        self.cv3 = TFConv(2 * c_, c2, 1, 1, w=w.cv3)
        self.m = keras.Sequential([TFBottleneck(c_, c_, shortcut, g, e=1.0,
            w=w.m[j]) for j in range(n)])

    def call(self, inputs):
        return self.cv3(tf.concat((self.m(self.cv1(inputs)), self.cv2(
            inputs)), axis=3))


class TFSPP(keras.layers.Layer):

    def __init__(self, c1, c2, k=(5, 9, 13), w=None):
        super(TFSPP, self).__init__()
        c_ = c1 // 2
        self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
        self.cv2 = TFConv(c_ * (len(k) + 1), c2, 1, 1, w=w.cv2)
        self.m = [keras.layers.MaxPool2D(pool_size=x, strides=1, padding=
            'SAME') for x in k]

    def call(self, inputs):
        x = self.cv1(inputs)
        return self.cv2(tf.concat([x] + [m(x) for m in self.m], 3))


class TFSPPF(keras.layers.Layer):

    def __init__(self, c1, c2, k=5, w=None):
        super(TFSPPF, self).__init__()
        c_ = c1 // 2
        self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1)
        self.cv2 = TFConv(c_ * 4, c2, 1, 1, w=w.cv2)
        self.m = keras.layers.MaxPool2D(pool_size=k, strides=1, padding='SAME')

    def call(self, inputs):
        x = self.cv1(inputs)
        y1 = self.m(x)
        y2 = self.m(y1)
        return self.cv2(tf.concat([x, y1, y2, self.m(y2)], 3))


class TFDetect(keras.layers.Layer):

    def __init__(self, nc=80, anchors=(), ch=(), imgsz=(640, 640), w=None):
        super(TFDetect, self).__init__()
        self.stride = tf.convert_to_tensor(w.stride.numpy(), dtype=tf.float32)
        self.nc = nc
        self.no = nc + 5
        self.nl = len(anchors)
        self.na = len(anchors[0]) // 2
        self.grid = [tf.zeros(1)] * self.nl
        self.anchors = tf.convert_to_tensor(w.anchors.numpy(), dtype=tf.float32
            )
        self.anchor_grid = tf.reshape(self.anchors * tf.reshape(self.stride,
            [self.nl, 1, 1]), [self.nl, 1, -1, 1, 2])
        self.m = [TFConv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in
            enumerate(ch)]
        self.training = False
        self.imgsz = imgsz
        for i in range(self.nl):
            ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1
                ] // self.stride[i]
            self.grid[i] = self._make_grid(nx, ny)

    def call(self, inputs):
        z = []
        x = []
        for i in range(self.nl):
            x.append(self.m[i](inputs[i]))
            ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1
                ] // self.stride[i]
            x[i] = tf.transpose(tf.reshape(x[i], [-1, ny * nx, self.na,
                self.no]), [0, 2, 1, 3])
            if not self.training:
                y = tf.sigmoid(x[i])
                xy = (y[..., 0:2] * 2.0 - 0.5 + self.grid[i]) * self.stride[i]
                wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]
                xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=
                    tf.float32)
                wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=
                    tf.float32)
                y = tf.concat([xy, wh, y[..., 4:]], -1)
                z.append(tf.reshape(y, [-1, 3 * ny * nx, self.no]))
        return x if self.training else (tf.concat(z, 1), x)

    @staticmethod
    def _make_grid(nx=20, ny=20):
        xv, yv = tf.meshgrid(tf.range(nx), tf.range(ny))
        return tf.cast(tf.reshape(tf.stack([xv, yv], 2), [1, 1, ny * nx, 2]
            ), dtype=tf.float32)


class TFUpsample(keras.layers.Layer):

    def __init__(self, size, scale_factor, mode, w=None):
        super(TFUpsample, self).__init__()
        assert scale_factor == 2, 'scale_factor must be 2'
        self.upsample = lambda x: tf.image.resize(x, (tuple(x.shape)[1] * 2,
            tuple(x.shape)[2] * 2), method=mode)

    def call(self, inputs):
        return self.upsample(inputs)


class TFConcat(keras.layers.Layer):

    def __init__(self, dimension=1, w=None):
        super(TFConcat, self).__init__()
        assert dimension == 1, 'convert only NCHW to NHWC concat'
        self.d = 3

    def call(self, inputs):
        return tf.concat(inputs, self.d)


def parse_model(d, ch, model, imgsz):
    LOGGER.info('\n%3s%18s%3s%10s  %-40s%-30s' % ('', 'from', 'n', 'params',
        'module', 'arguments'))
    anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d[
        'width_multiple']
    na = len(anchors[0]) // 2 if isinstance(anchors, list) else anchors
    no = na * (nc + 5)
    layers, save, c2 = [], [], ch[-1]
    for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):
        m_str = m
        m = eval(m) if isinstance(m, str) else m
        for j, a in enumerate(args):
            try:
                args[j] = eval(a) if isinstance(a, str) else a
            except NameError:
                pass
        n = max(round(n * gd), 1) if n > 1 else n
        if m in [paddle.nn.Conv2D, Conv, Bottleneck, SPP, SPPF, DWConv,
            MixConv2d, Focus, CrossConv, BottleneckCSP, C3]:
            c1, c2 = ch[f], args[0]
            c2 = make_divisible(c2 * gw, 8) if c2 != no else c2
            args = [c1, c2, *args[1:]]
            if m in [BottleneckCSP, C3]:
                args.insert(2, n)
                n = 1
        elif m is paddle.nn.BatchNorm2D:
            args = [ch[f]]
        elif m is Concat:
            c2 = sum([ch[-1 if x == -1 else x + 1] for x in f])
        elif m is Detect:
            args.append([ch[x + 1] for x in f])
            if isinstance(args[1], int):
                args[1] = [list(range(args[1] * 2))] * len(f)
            args.append(imgsz)
        else:
            c2 = ch[f]
        tf_m = eval('TF' + m_str.replace('nn.', ''))
        m_ = keras.Sequential([tf_m(*args, w=model.model[i][j]) for j in
            range(n)]) if n > 1 else tf_m(*args, w=model.model[i])
        torch_m_ = paddle.nn.Sequential(*[m(*args) for _ in range(n)]
            ) if n > 1 else m(*args)
        t = str(m)[8:-2].replace('__main__.', '')
        np = sum([x.size for x in torch_m_.parameters()])
        m_.i, m_.f, m_.type, m_.np = i, f, t, np
        LOGGER.info('%3s%18s%3s%10.0f  %-40s%-30s' % (i, f, n, np, t, args))
        save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x !=
            -1)
        layers.append(m_)
        ch.append(c2)
    return keras.Sequential(layers), sorted(save)


class TFModel:

    def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, model=None, imgsz
        =(640, 640)):
        super(TFModel, self).__init__()
        if isinstance(cfg, dict):
            self.yaml = cfg
        else:
            import yaml
            self.yaml_file = Path(cfg).name
            with open(cfg) as f:
                self.yaml = yaml.load(f, Loader=yaml.FullLoader)
        if nc and nc != self.yaml['nc']:
            print('Overriding %s nc=%g with nc=%g' % (cfg, self.yaml['nc'], nc)
                )
            self.yaml['nc'] = nc
        self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch
            ], model=model, imgsz=imgsz)

    def predict(self, inputs, tf_nms=False, agnostic_nms=False,
        topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25):
        y = []
        x = inputs
        for i, m in enumerate(self.model.layers):
            if m.f != -1:
                x = y[m.f] if isinstance(m.f, int) else [(x if j == -1 else
                    y[j]) for j in m.f]
            x = m(x)
            y.append(x if m.i in self.savelist else None)
        if tf_nms:
            boxes = self._xywh2xyxy(x[0][..., :4])
            probs = x[0][:, :, 4:5]
            classes = x[0][:, :, 5:]
            scores = probs * classes
            if agnostic_nms:
                nms = AgnosticNMS()((boxes, classes, scores), topk_all,
                    iou_thres, conf_thres)
                return nms, x[1]
            else:
                boxes = tf.expand_dims(boxes, 2)
                nms = tf.image.combined_non_max_suppression(boxes, scores,
                    topk_per_class, topk_all, iou_thres, conf_thres,
                    clip_boxes=False)
                return nms, x[1]
        return x[0]

    @staticmethod
    def _xywh2xyxy(xywh):
        x, y, w, h = tf.split(xywh, num_or_size_splits=4, axis=-1)
        return tf.concat([x - w / 2, y - h / 2, x + w / 2, y + h / 2], axis=-1)


class AgnosticNMS(keras.layers.Layer):

    def call(self, input, topk_all, iou_thres, conf_thres):
        return tf.map_fn(lambda x: self._nms(x, topk_all, iou_thres,
            conf_thres), input, fn_output_signature=(tf.float32, tf.float32,
            tf.float32, tf.int32), name='agnostic_nms')

    @staticmethod
    def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25):
        boxes, classes, scores = x
        class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32)
        scores_inp = tf.reduce_max(scores, -1)
        selected_inds = tf.image.non_max_suppression(boxes, scores_inp,
            max_output_size=topk_all, iou_threshold=iou_thres,
            score_threshold=conf_thres)
        selected_boxes = tf.gather(boxes, selected_inds)
        padded_boxes = tf.pad(selected_boxes, paddings=[[0, topk_all - tf.
            shape(selected_boxes)[0]], [0, 0]], mode='CONSTANT',
            constant_values=0.0)
        selected_scores = tf.gather(scores_inp, selected_inds)
        padded_scores = tf.pad(selected_scores, paddings=[[0, topk_all - tf
            .shape(selected_boxes)[0]]], mode='CONSTANT', constant_values=-1.0)
        selected_classes = tf.gather(class_inds, selected_inds)
        padded_classes = tf.pad(selected_classes, paddings=[[0, topk_all -
            tf.shape(selected_boxes)[0]]], mode='CONSTANT', constant_values
            =-1.0)
        valid_detections = tf.shape(selected_inds)[0]
        return padded_boxes, padded_scores, padded_classes, valid_detections


def representative_dataset_gen(dataset, ncalib=100):
    for n, (path, img, im0s, vid_cap) in enumerate(dataset):
        input = np.transpose(img, [1, 2, 0])
        input = np.expand_dims(input, axis=0).astype(np.float32)
        input /= 255.0
        yield [input]
        if n >= ncalib:
            break


def run(weights=ROOT / 'yolov5s.pt', imgsz=(640, 640), batch_size=1,
    dynamic=False):
    im = paddle.zeros(shape=(batch_size, 3, *imgsz))
    model = attempt_load(weights, map_location=paddle.CPUPlace(), inplace=
        True, fuse=False)
    y = model(im)
    model.info()
    im = tf.zeros((batch_size, *imgsz, 3))
    tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
    y = tf_model.predict(im)
    im = keras.Input(shape=(*imgsz, 3), batch_size=None if dynamic else
        batch_size)
    keras_model = keras.Model(inputs=im, outputs=tf_model.predict(im))
    keras_model.summary()


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt',
        help='weights path')
    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=
        int, default=[640], help='inference size h,w')
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    parser.add_argument('--dynamic', action='store_true', help=
        'dynamic batch size')
    opt = parser.parse_args()
    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1
    print_args(FILE.stem, opt)
    return opt


def main(opt):
    set_logging()
    run(**vars(opt))


if __name__ == '__main__':
    opt = parse_opt()
    main(opt)
